#coding=utf8
# Copyright (c) 2016 Tinydot. inc.
# All Rights Reserved.
#
#    Licensed under the Apache License, Version 2.0 (the "License"); you may
#    not use this file except in compliance with the License. You may obtain
#    a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#    License for the specific language governing permissions and limitations
#    under the License.


import numpy as np
from PIL import Image
from skimage.measure import find_contours
from torch.cuda import is_available as cuda_is_available
from scipy.ndimage.measurements import label as np_label
from scipy.ndimage.morphology import binary_dilation
import torch


if cuda_is_available():
    import cupyx.scipy.ndimage.label as cp_label


def get_instance_contours(instance_image, contour_value=None, force_return_np=False, iterations=0):
    instance = instance_image
    r_mask = np.zeros_like(instance_image)
    if not isinstance(instance_image, np.ndarray):
        instance = np.array(instance_image)
        contour_value = contour_value or 255
    else:
        contour_value = contour_value or 1.0
    cts = []
    for u in np.unique(instance):
        if u == 0:
            continue
        cts.extend(find_contours(instance == u))
    # return instance_image
    for ct in cts:
        r_mask[np.round(ct[:, 0]).astype('int'), np.round(ct[:, 1]).astype('int')] = contour_value
    if iterations > 0:
        r_mask = binary_dilation(r_mask, iterations=1)
    if isinstance(instance_image, np.ndarray) or force_return_np:
        return r_mask
    return Image.fromarray(r_mask.astype(np.uint8))


def contours_as_label(input_contours_map, thresh_hold):
    contours_map = input_contours_map
    # contours_map: w, h
    return_img = False
    if isinstance(contours_map, Image.Image):
        contours_map = np.array(contours_map)
        return_img = True

    if isinstance(contours_map, np.ndarray):
        cts_map = (contours_map < thresh_hold).astype(np.int)
        labels, _ = np_label(cts_map)
        cts_map = cts_map == 0
        if return_img:
            from .uniquecolors import ins_numpy_to_dis_image
            return ins_numpy_to_dis_image(labels), cts_map
        return labels, cts_map

    if isinstance(contours_map, torch.Tensor):
        if not cuda_is_available():
            l, m = contours_as_label(contours_map.numpy(), thresh_hold)
            return torch.from_numpy(l).to(contours_map.device), torch.from_numpy(m).to(contours_map.device)
        else:
            cts_map = (contours_map < thresh_hold).int()
            labels, _ = cp_label(cts_map)
            cts_map = cts_map == 0
            return labels, cts_map.int()

    assert False, "data type [%s] not supported" % type(contours_map)
